package org.gwh.airagknowledge.util;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Profile;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.jdbc.core.JdbcTemplate;

@Slf4j
@Configuration
public class DatabaseConnectionTester {

    @Bean
    @Profile("!test")
    public CommandLineRunner testDatabaseConnections(
            @Autowired JdbcTemplate jdbcTemplate,
            @Autowired(required = false) RedisTemplate<String, Object> redisTemplate) {
        
        return args -> {
            // 测试PostgreSQL连接
            try {
                log.info("Testing PostgreSQL connection...");
                String postgresVersion = jdbcTemplate.queryForObject("SELECT version()", String.class);
                log.info("PostgreSQL connection successful! Version: {}", postgresVersion);
                
                // 测试pgvector扩展
                try {
                    Boolean vectorExtensionExists = jdbcTemplate.queryForObject(
                            "SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'vector')", 
                            Boolean.class);
                    
                    if (Boolean.TRUE.equals(vectorExtensionExists)) {
                        log.info("pgvector extension is installed!");
                    } else {
                        log.warn("pgvector extension is NOT installed! Vector operations will not work.");
                    }
                } catch (Exception e) {
                    log.error("Error checking pgvector extension: {}", e.getMessage());
                }
                
            } catch (Exception e) {
                log.error("PostgreSQL connection failed: {}", e.getMessage());
            }
            
            // 测试Redis连接
            if (redisTemplate != null) {
                try {
                    log.info("Testing Redis connection...");
                    String testKey = "connection_test";
                    String testValue = "connected_at_" + System.currentTimeMillis();
                    
                    redisTemplate.opsForValue().set(testKey, testValue);
                    Object retrievedValue = redisTemplate.opsForValue().get(testKey);
                    redisTemplate.delete(testKey);
                    
                    if (testValue.equals(retrievedValue)) {
                        log.info("Redis connection successful!");
                    } else {
                        log.warn("Redis connection issue: test value doesn't match.");
                    }
                } catch (Exception e) {
                    log.error("Redis connection failed: {}", e.getMessage());
                }
            } else {
                log.warn("Redis template not available. Skipping Redis connection test.");
            }
        };
    }
} 